#!/bin/env python
# -*- coding: utf-8 -*-


"""
File Name: dump_hbase_meta.py
Author: morre <morre@vip.qq.com>
Create Time: 2022/12/31 17:35:16
Brief:
"""
import io
import sys
import json
import base64
import logging
import argparse
from hbase_client import HbaseClient


def main():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s  %(filename)s[line:%(lineno)d] %(levelname)s "
        "%(message)s",
        datefmt="%a, %d %b %Y %H:%M:%S")

    parser = argparse.ArgumentParser(
        description="dump hbase meta")

    parser.add_argument("--hbase_conf", dest="hbase_conf_file",
                        default="hbase.conf",
                        help="hbase config file")

    parser.add_argument("--table_name", dest="table_name",
                        help="hbase table_name")

    parser.add_argument("--db_path", dest="db_path",
                        help="db path")

    parser.add_argument(
        "--action", dest="action", default="dump_hbase_meta",
        help="dump_hbase_meta,find_region_server,get_hbase_meta")

    parser.add_argument(
        "--base64_encode", dest="base64_encode", action="store_true",
        help="base64 encode rowkey")

    options = parser.parse_args(sys.argv[1:])

    obj = {}
    with open(options.hbase_conf_file, "r", encoding="utf8") as fp:
        obj = json.load(fp)

    client = HbaseClient(
        zkquorum=obj["zkquorum"], zk_root_path=obj["zk_root_path"],
        init_socket_pool_size=1, establish_connection_timeout=60,
        missing_znode_retries=3, rpc_socket_timeout=60,
        max_write_buff_size=1000)
    client.open()

    if options.action == "dump_hbase_meta":
        client.dump_meta_to_sqlite(options.db_path, options.table_name.encode())
    elif options.action in ["find_region_server", "find_region_name"]:
        if options.db_path:
            client.load_meta_from_sqlite(options.db_path)
        input_fifo = io.TextIOWrapper(
            sys.stdin.buffer, encoding="utf8", errors="ignore")
        for line in input_fifo:
            fields = line.strip("\n").split("\t")
            region_info = client.find_region(
                options.table_name.encode(), fields[0].encode())
            res = b"null"
            if region_info:
                if options.action == "find_region_server":
                    res = region_info.server_loc
                elif options.action == "find_region_name":
                    res = region_info.region_name
            print("\t".join([res.decode()] + fields))
    elif options.action == "get_hbase_meta":
        lines = set()
        for table, region_name, start_key, stop_key, server_loc in \
                client.get_table_regions(options.table_name.encode(), 20000):
            if options.base64_encode:
                region_name = base64.b64encode(region_name)
                start_key = base64.b64encode(start_key)
                stop_key = base64.b64encode(stop_key)
            line = b"\t".join(
                [table, region_name, start_key, stop_key, server_loc]
            ).decode("utf8")
            if line in lines:
                continue
            print(line)
            lines.add(line)
    return 0


if __name__ == "__main__":
    sys.exit(main())
